#!/usr/bin/env python3
"""
RAG问答系统 - 检索增强生成

这个模块实现了一个完整的RAG（检索增强生成）问答系统，结合了：
1. 向量检索：从知识库中检索相关文档
2. 大模型生成：使用ModelArts API生成基于知识的回答

功能特性：
- 支持中文问答
- 自动从向量数据库检索相关文档
- 使用盘古大模型生成准确回答
- 提供检索结果的来源信息
- 支持自定义检索参数

作者：AI Assistant
日期：2025-09-27
版本：1.0.0
"""

import os
import sys
import logging
from typing import List, Dict, Optional, Tuple
from pathlib import Path
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 导入OpenAI客户端和相关模块
from openai import OpenAI

# HuggingFace嵌入模型（已弃用，改用API调用）
# from langchain_community.embeddings import HuggingFaceEmbeddings
# BGE-M3 API Embeddings
from bge_m3_api_embeddings import BGEM3APIEmbeddings

# 导入Langchain相关模块
from langchain.schema import Document
from langchain_community.vectorstores import Chroma

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class RAGQASystem:
    """RAG问答系统主类"""
    
    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.modelarts-maas.com/v1",
        model_name: str = "openpangu-pro-moe-72b",
        vector_db_path: str = "./chroma_db",
        embedding_model_path: str = None  # 已弃用，保留以保持兼容性
    ):
        """
        初始化RAG问答系统
        
        Args:
            api_key: ModelArts API密钥
            base_url: API基础URL
            model_name: 使用的大模型名称
            vector_db_path: 向量数据库路径
            embedding_model_path: 嵌入模型本地路径（已弃用，现在使用API调用）
        """
        self.api_key = api_key
        self.base_url = base_url
        self.model_name = model_name
        self.vector_db_path = vector_db_path
        self.embedding_model_path = embedding_model_path  # 保留以保持兼容性
        
        # 初始化OpenAI客户端
        self.client = OpenAI(api_key=api_key, base_url=base_url)
        
        # 初始化嵌入模型和向量数据库
        self.embeddings_model = None
        self.vector_store = None
        
        logger.info("RAG问答系统初始化完成")
    
    def _get_optimal_device(self) -> str:
        """获取最优设备配置"""
        try:
            import torch
            
            if torch.cuda.is_available():
                try:
                    memory_total = torch.cuda.get_device_properties(0).total_memory / 1024**2
                    memory_allocated = torch.cuda.memory_allocated() / 1024**2
                    memory_free = memory_total - memory_allocated
                    
                    if memory_free > 2048:
                        device_name = torch.cuda.get_device_name()
                        logger.info(f"🚀 使用GPU: {device_name}")
                        return 'cuda'
                    else:
                        logger.warning(f"⚠️  GPU内存不足，使用CPU")
                except Exception as e:
                    logger.warning(f"⚠️  GPU检查失败: {e}，使用CPU")
            
            if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
                logger.info("🍎 使用Apple Silicon MPS")
                return 'mps'
            
            logger.info("💻 使用CPU")
            return 'cpu'
            
        except ImportError:
            logger.warning("⚠️  PyTorch未安装，使用CPU")
            return 'cpu'
    
    def load_vector_database(self) -> bool:
        """
        加载向量数据库和嵌入模型
        
        Returns:
            bool: 是否加载成功
        """
        try:
            # 检查向量数据库是否存在
            if not os.path.exists(self.vector_db_path):
                logger.error(f"向量数据库不存在: {self.vector_db_path}")
                logger.error("请先运行 rag_document_ingestion.py 来创建向量数据库")
                return False
            
            # 初始化嵌入模型（使用API调用）
            logger.info("初始化BGE-M3 API嵌入模型...")
            self.embeddings_model = BGEM3APIEmbeddings()
            logger.info("✅ BGE-M3 API嵌入模型初始化成功")
            
            # 加载向量数据库
            logger.info("加载向量数据库...")
            self.vector_store = Chroma(
                persist_directory=self.vector_db_path,
                embedding_function=self.embeddings_model,
                collection_name="rag_documents"
            )
            
            # 验证数据库内容
            collection = self.vector_store._collection
            count = collection.count()
            logger.info(f"✅ 向量数据库加载成功，包含 {count} 个文档片段")
            
            return True
            
        except Exception as e:
            logger.error(f"❌ 加载向量数据库失败: {str(e)}")
            return False
    
    def search_relevant_documents(
        self,
        query: str,
        k: int = 5,
        score_threshold: float = 1.0
    ) -> List[Tuple[Document, float]]:
        """
        搜索相关文档
        
        Args:
            query: 查询问题
            k: 返回的文档数量
            score_threshold: 距离阈值（ChromaDB返回距离值，越小越相似，建议值：0.5-1.5）
            
        Returns:
            List[Tuple[Document, float]]: 文档和距离值的列表
        """
        if not self.vector_store:
            logger.error("向量数据库未加载")
            return []
        
        try:
            logger.info(f"🔍 检索相关文档: '{query}'")
            
            # 执行相似度搜索（返回的是距离值，越小越相似）
            results_with_scores = self.vector_store.similarity_search_with_score(
                query, k=k
            )
            
            # 过滤距离过大的结果（距离越小越相似，所以保留 <= threshold 的结果）
            filtered_results = [
                (doc, score) for doc, score in results_with_scores
                if score <= score_threshold
            ]
            
            logger.info(f"找到 {len(filtered_results)} 个相关文档片段（距离阈值: {score_threshold}）")
            
            return filtered_results
            
        except Exception as e:
            logger.error(f"❌ 文档检索失败: {str(e)}")
            return []
    
    def create_context_prompt(
        self,
        query: str,
        relevant_docs: List[Tuple[Document, float]]
    ) -> str:
        """
        创建包含上下文的提示词
        
        Args:
            query: 用户问题
            relevant_docs: 相关文档列表（包含距离值）
            
        Returns:
            str: 完整的提示词
        """
        if not relevant_docs:
            return f"""请回答以下问题：

问题：{query}

注意：由于没有找到相关的知识库内容，请基于你的通用知识回答，并说明这个回答不是基于特定知识库的。"""
        
        # 构建上下文
        context_parts = []
        for i, (doc, distance) in enumerate(relevant_docs):
            source_file = doc.metadata.get('source_file', '未知来源')
            source_name = os.path.basename(source_file) if source_file != '未知来源' else '未知来源'
            
            # 将距离值转换为相似度百分比（距离越小，相似度越高）
            # 使用简单的转换公式：similarity = max(0, 1 - distance)
            similarity = max(0, 1 - distance) * 100
            
            context_parts.append(
                f"【参考资料{i+1}】（来源：{source_name}，相似度：{similarity:.1f}%）\n{doc.page_content}\n"
            )
        
        context = "\n".join(context_parts)
        
        # 创建完整提示词
        prompt = f"""请基于以下参考资料回答用户问题。请确保回答准确、详细且有逻辑性。如果参考资料不足以完全回答问题，请明确说明并基于可用信息给出部分回答。

参考资料：
{context}

问题：{query}

请按以下格式回答：
1. 先给出简洁的答案
2. 然后提供详细解释
3. 最后列出主要参考来源

回答："""
        
        return prompt
    
    def call_llm_api(self, prompt: str) -> Optional[str]:
        """
        调用大模型API
        
        Args:
            prompt: 提示词
            
        Returns:
            Optional[str]: 模型回答，失败时返回None
        """
        try:
            logger.info("🤖 调用大模型API生成回答...")
            
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {
                        "role": "system", 
                        "content": "你是一个专业的AI助手，擅长基于提供的参考资料回答问题。你会仔细阅读参考资料，给出准确、有条理的回答，并明确标注信息来源。"
                    },
                    {
                        "role": "user", 
                        "content": prompt
                    }
                ],
                temperature=0.1,  # 降低随机性，提高回答的一致性
                max_tokens=2000   # 限制回答长度
            )
            
            answer = response.choices[0].message.content
            logger.info("✅ 大模型回答生成完成")
            
            return answer
            
        except Exception as e:
            logger.error(f"❌ 调用大模型API失败: {str(e)}")
            return None
    
    def ask_question(
        self,
        question: str,
        max_docs: int = 5,
        score_threshold: float = 1.0
    ) -> Dict[str, any]:
        """
        问答主函数
        
        Args:
            question: 用户问题
            max_docs: 最大检索文档数
            score_threshold: 距离阈值（ChromaDB返回距离值，越小越相似，建议值：0.5-1.5）
            
        Returns:
            Dict: 包含答案和元信息的字典
        """
        logger.info(f"📝 收到问题: {question}")
        
        result = {
            "question": question,
            "answer": None,
            "sources": [],
            "retrieved_docs": 0,
            "error": None
        }
        
        try:
            # 检查系统是否已初始化
            if not self.vector_store:
                logger.error("系统未初始化，请先调用 load_vector_database()")
                result["error"] = "系统未初始化"
                return result
            
            # 1. 检索相关文档
            relevant_docs = self.search_relevant_documents(
                question, k=max_docs, score_threshold=score_threshold
            )
            result["retrieved_docs"] = len(relevant_docs)
            
            # 记录来源信息
            sources = []
            for doc, distance in relevant_docs:
                source_file = doc.metadata.get('source_file', '未知来源')
                source_name = os.path.basename(source_file) if source_file != '未知来源' else '未知来源'
                # 将距离值转换为相似度百分比（距离越小，相似度越高）
                similarity = max(0, 1 - distance) * 100
                sources.append({
                    "file": source_name,
                    "distance": distance,  # 保留原始距离值
                    "similarity": similarity,  # 添加相似度百分比
                    "preview": doc.page_content[:100] + "..." if len(doc.page_content) > 100 else doc.page_content
                })
            result["sources"] = sources
            
            # 2. 创建提示词
            prompt = self.create_context_prompt(question, relevant_docs)
            
            # 3. 调用大模型生成回答
            answer = self.call_llm_api(prompt)
            
            if answer:
                result["answer"] = answer
                logger.info("✅ 问答完成")
            else:
                result["error"] = "大模型API调用失败"
                logger.error("❌ 大模型API调用失败")
            
            return result
            
        except Exception as e:
            error_msg = f"问答过程出错: {str(e)}"
            logger.error(f"❌ {error_msg}")
            result["error"] = error_msg
            return result
    
    def interactive_qa(self):
        """交互式问答模式"""
        print("🤖 RAG问答系统已启动")
        print("=" * 60)
        print("💡 使用说明：")
        print("   - 直接输入问题即可获得基于知识库的回答")
        print("   - 输入 'quit' 或 'exit' 退出")
        print("   - 输入 'help' 查看更多命令")
        print("=" * 60)
        
        while True:
            try:
                # 获取用户输入
                question = input("\n🙋 请输入您的问题: ").strip()
                
                # 处理特殊命令
                if question.lower() in ['quit', 'exit', '退出']:
                    print("👋 再见！")
                    break
                elif question.lower() == 'help':
                    print("\n📖 可用命令：")
                    print("   help  - 显示帮助信息")
                    print("   quit  - 退出系统")
                    print("   exit  - 退出系统")
                    continue
                elif not question:
                    print("❓ 请输入一个问题")
                    continue
                
                # 处理问题
                result = self.ask_question(question)
                
                # 显示结果
                print(f"\n📊 检索统计: 找到 {result['retrieved_docs']} 个相关文档片段")
                
                if result['sources']:
                    print(f"📚 主要来源:")
                    for i, source in enumerate(result['sources'][:3]):  # 只显示前3个来源
                        print(f"   {i+1}. {source['file']} (相似度: {source['similarity']:.1f}%)")
                
                if result['answer']:
                    print(f"\n🤖 回答:")
                    print("-" * 50)
                    print(result['answer'])
                    print("-" * 50)
                elif result['error']:
                    print(f"\n❌ 错误: {result['error']}")
                
            except KeyboardInterrupt:
                print("\n\n👋 用户中断，系统退出")
                break
            except Exception as e:
                print(f"\n❌ 发生错误: {str(e)}")


def main():
    """主函数"""
    print("🚀 RAG问答系统")
    print("=" * 50)
    
    # 配置信息 - 从环境变量读取
    API_KEY = os.getenv("MODELARTS_API_KEY")
    BASE_URL = os.getenv("MODELARTS_BASE_URL", "https://api.modelarts-maas.com/v1")
    MODEL_NAME = os.getenv("MODELARTS_MODEL_NAME", "openpangu-pro-moe-72b")
    
    # 验证必需的环境变量
    if not API_KEY:
        print("❌ 错误：未设置 MODELARTS_API_KEY 环境变量")
        print("请创建 .env 文件并设置 MODELARTS_API_KEY，参考 .env.example 文件")
        return False
    
    # 创建RAG系统实例
    rag_system = RAGQASystem(
        api_key=API_KEY,
        base_url=BASE_URL,
        model_name=MODEL_NAME
    )
    
    # 加载向量数据库
    print("🔄 加载向量数据库和嵌入模型...")
    if not rag_system.load_vector_database():
        print("❌ 初始化失败！")
        print("\n🔧 解决方案：")
        print("1. 确保已运行 python rag_document_ingestion.py 创建向量数据库")
        print("2. 检查 ./chroma_db 目录是否存在")
        print("3. 检查 API 密钥是否已配置（BGE_M3_API_KEY 或 MAAS_API_KEY）")
        return False
    
    print("✅ 系统初始化完成！")
    
    # 启动交互式问答
    try:
        rag_system.interactive_qa()
        return True
    except Exception as e:
        print(f"❌ 系统运行出错: {str(e)}")
        return False


if __name__ == "__main__":
    """
    RAG问答系统主程序
    
    使用前准备：
    1. 确保已安装必要依赖：
       pip install openai langchain langchain-community chromadb requests
    
    2. 准备向量数据库：
       python rag_document_ingestion.py
    
    3. 配置 API 密钥（BGE_M3_API_KEY 或 MAAS_API_KEY 环境变量）
    
    4. 运行系统：
       python rag_qa_system.py
    """
    success = main()
    sys.exit(0 if success else 1)
